import unittest, traceback, os, tempfile
from ConfigParser import ConfigParser

from pdboi.pdbents.entry import Entry
from pdboi.pdbdata.xmldataconnector import XMLDataConnector
from pdboi.pdbdata.mmcifapiconnector import MMcifAPIconnector
from pdboi.pdbents.propertystore import PropertyStore
from pdboi.pdbents.group import ChemicalGroup

#@unittest.skip("Skipping ConnectorTests")
class ConnectorTests(unittest.TestCase) :
	'''test that entry title and entity info for 1cbs as given by Entry object are identical despite use of various connectors'''
	def setUp(s) :
		s.connectors = [
			MMcifAPIconnector("http://www.ebi.ac.uk/~swanand/1cbs.cif.gz"),
			#XMLDataConnector("/ebi/ftp/pub/databases/pdb/data/structures/all/XML/1cbs.xml.gz"),
			#XMLDataConnector("http://www.ebi.ac.uk/pdbe-srv/view/files/1cbs.xml"),
			#XMLDataConnector("http://www.ebi.ac.uk/~swanand/1cbs.xml.gz"),
		]
		# from configs, read info about 1cbs: title; entity info - name, type, polytype, struct-asym-ids associated ; a listing of residue numbers
		cfg = ConfigParser()
		cfg.read("configs.txt")
		s.entryTitle = cfg.get("TestdataFor1cbs", "title")
		s.resnums = eval(cfg.get("TestdataFor1cbs", "resnums"))
		s.entityinfo = eval(cfg.get("TestdataFor1cbs", "entityinfo"))
	def testEntryEntityTitles(s) :
		for dc in s.connectors :
			entry = Entry(dc)
			s.assertEqual(s.entryTitle, entry.getTitle())
			ents = entry.getEntities()
			for ei in range(len(ents)) :
				s.assertEqual(ents[ei].getName(), s.entityinfo[ei][0])
				s.assertEqual(ents[ei].type, s.entityinfo[ei][1])
				s.assertEqual(ents[ei].polytype, s.entityinfo[ei][2])
				saids = [inst.getStructAsymId() for inst in ents[ei].getInstances()]
				s.assertEqual(saids, s.entityinfo[ei][3])
	def testIterationThroughPDBhierarchy(s) :
		for dc in s.connectors :
			entry = Entry(dc)
			modelChainResnums = []
			for model in entry.getModels() :
				for chain in model.getChains() :
					for group in chain.getModelledGroups() :
						modelChainResnums.append( (model.getModelId(), chain.getChainId(), group.getIcodedResnum()) )
			s.assertEqual(modelChainResnums, s.resnums)

#@unittest.skip("Skipping ChemicalgroupTest")
class EntityGroupModellingTest(unittest.TestCase) :
	'''test that modelled and unmodelled stretches of polymeric entities are reported properly'''
	def testModelledStretches(s) :
		dc = MMcifAPIconnector("http://www.ebi.ac.uk/~swanand/1aac.cif.gz")
		entry = Entry(dc)
		tempfn = tempfile.mkstemp()[1]
		ofp = open(tempfn,'w')
		for model in entry.getModels() :
			for chain in model.getChains() :
				for group in chain.getModelledGroups() :
					print >> ofp, model.getModelId(), chain.getChainId(), group.getIcodedResnum()
					subgroups = group.getSubgroups()
					if len(subgroups) > 1 : print >> ofp, [sg.getId() for sg in subgroups]
		ofp.close()
		lines1 = open(tempfn).readlines()
		lines2 = open("testModelledStretches.out").readlines()
		for li in range(len(lines1)) : s.assertEqual(lines1[li], lines2[li])
		#if group.getSubgroups() != '' :
		#	print model.getModelId(), chain.getChainId(), group.getIcodedResnum(), group.getSubgroups()

class ChemicalgroupTest(unittest.TestCase) :
	def testChemgroupCreation(s) :
		'''checks that chemical group ATP is created properly with expected number of bonds and atoms, and converted to openbabel molecule, and mol2 file is written as expected'''
		cg = ChemicalGroup.getChemicalGroup("ATP")
		s.assertEqual(cg.getName(), "ATP")
		s.assertEqual(cg.numAtoms(),47)
		s.assertEqual(cg.numBonds(),49)
		obmol = cg.makeOpenbabelMolecule()
		atomnames = []
		for aind in range(1,obmol.NumAtoms()+1) :
			atomnames.append(obmol.GetResidue(0).GetAtomID( obmol.GetAtom(aind) ))
		s.assertEqual( len(atomnames), 47 )
		idealAtomnames = ["PG", "O1G", "O2G", "O3G", "PB", "O1B", "O2B", "O3B", "PA", "O1A", "O2A", "O3A", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "O2'", "C1'", "N9", "C8", "N7", "C5", "C6", "N6", "N1", "C2", "N3", "C4", "HOG2", "HOG3", "HOB2", "HOA2", "H5'1", "H5'2", "H4'", "H3'", "HO3'", "H2'", "HO2'", "H1'", "H8", "HN61", "HN62", "H2",]
		s.assertEqual( idealAtomnames, atomnames )
		bondatoms = []
		arombonds = [ ("N9","C8",), ("N9","C4",), ("C8","N7",), ("N7","C5",), ("C5","C6",), ("C5","C4",), ("C6","N1",), ("N1","C2",), ("C2","N3",), ("N3","C4",), ]
		for bind in range(obmol.NumBonds()) :
			obbond = obmol.GetBond(bind)
			atompair = ( (obmol.GetResidue(0).GetAtomID( obbond.GetBeginAtom() ), obmol.GetResidue(0).GetAtomID( obbond.GetEndAtom() )) )
			bondatoms.append( atompair )
			if atompair in arombonds or tuple(reversed(atompair)) in arombonds :
				s.assertEqual(True, obbond.IsAromatic()) # bond order is returned as 1 or 2 though....
		s.assertEqual( len(bondatoms), 49 )
		atompairs = []
		for abond in cg.getBonds() : atompairs.append( (abond.getAtomId1(), abond.getAtomId2()) )
		s.assertEqual( len(atompairs), 49 )
		for bo in bondatoms : s.assertIn( bo, atompairs )
		for bo in atompairs : s.assertIn( bo, bondatoms )
		## now check that atp.mol2 is written properly
		entry = Entry( MMcifAPIconnector("http://www.ebi.ac.uk/pdbe-srv/view/files/1atp.cif") )
		for group in entry.getNonpolymericNonwaterGroups() :
			for subgr in group.getSubgroups() :
				if subgr.getResname() == "ATP" :
					mol2file = tempfile.mkstemp(suffix=".mol2")[1]
					cg.writeMol2file(subgr,mol2file)
					s.assertEqual(open("atp.mol2").readlines(), open(mol2file).readlines())
					os.unlink(mol2file)

class PropertyStoreTests(unittest.TestCase) :
	def testbasic(s) :
		s.assertEqual(1,1)
		dc = MMcifAPIconnector("http://www.ebi.ac.uk/~swanand/1aac.cif.gz")
		entry = Entry(dc)
		ps = PropertyStore.instance()
		ps.set(entry, "entryprop", "an-entry-property", "mysource")
		for model in entry.getModels() :
			ps = PropertyStore.instance()
			ps.set(model, "modelprop", "a-model-property", "mysource")
			for chain in model.getChains() :
				ps.set(chain, "chainprop", "a-chain-property", "mysource")
				for group in chain.getModelledGroups() :
					ps.set(group, "groupprop", "a-group-property", "mysource")
					for sg in group.getSubgroups() :
						ps.set(sg, "subgroupprop", "a-subgroup-property", "mysource")
		# check that another instance is the same to check singleton-ness
		ps1 = PropertyStore.instance()
		s.assertEqual(ps,ps1)
		# check some values set previously
		s.assertEqual(ps.get(entry,"entryprop"), "an-entry-property")
		for model in entry.getModels() :
			s.assertEqual(ps.get(model,"modelprop"),"a-model-property")
		# pickle to a file and load from it and check values set previously
		pklfile = tempfile.mkstemp()[1]
		ps.store(pklfile)
		ps2 = PropertyStore.instance()
		ps2.load(pklfile)
		s.assertEqual(ps2.get(entry,"entryprop"), "an-entry-property")
		for model in entry.getModels() :
			s.assertEqual(ps2.get(model,"modelprop"),"a-model-property")
		os.unlink(pklfile)

if __name__ == "__main__" : unittest.main()
